package com.cy.duosearch.langchain;

import com.alibaba.fastjson2.JSON;
import dev.langchain4j.model.chat.ChatLanguageModel;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;

import static com.cy.duosearch.langchain.PromptConstant.systemPrompt;

@RequestMapping("api")
@RestController
public class ChatController {

    ChatLanguageModel chatLanguageModel;

    public ChatController(ChatLanguageModel chatLanguageModel) {
        this.chatLanguageModel = chatLanguageModel;
    }

    @GetMapping("/chat")
    public String model(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        return chatLanguageModel.chat(message);
    }


    @Autowired
    Assistant assistant;

    @GetMapping("/chat2")
    public String chat(String message) {
        return assistant.chat(message);
    }

    public static Integer count = 0;
    @RequestMapping("process")
    public void process() {
        String folder = "data";
        extractPDFContents(folder);
    }

    /**
     * 遍历指定路径下的所有PDF文件，并提取内容存放到列表中
     * @param directoryPath 要遍历的目录路径
     * @return 包含每个PDF文件内容的列表
     */
    public void extractPDFContents(String directoryPath) {

        count = 0;

        File directory = new File(directoryPath);

        // 遍历目录下的所有文件和子目录
        File[] files = directory.listFiles();
        if (files != null) {
            for (File file : files) {
                // 如果是文件且扩展名是.pdf，则提取内容
                if (file.isFile() && file.getName().toLowerCase().endsWith(".pdf")) {
                    extractPDFContent(file.getAbsolutePath(),file.getName());
                    count++;
                    System.out.println("总量：" + count);
                }
                // 如果是目录，则递归遍历
                else if (file.isDirectory()) {
                    extractPDFContents(file.getAbsolutePath());
                }
            }
        }
    }


    /**
     * 提取单个PDF文件的内容
     * @param filePath PDF文件的路径
     * @return PDF文件的内容
     */
    private void extractPDFContent(String filePath,String fileName) {
        try (PDDocument document = PDDocument.load(new File(filePath))) {
            PDFTextStripper stripper = new PDFTextStripper();
            String content = stripper.getText(document);
            String requestPrompt = systemPrompt + "下面的安全报告的内容：" + content;

            // 创建用户消息

//            String aiAnswer = chatLanguageModel.chat(requestPrompt);
            System.out.println("写入json文件");

            // 调用 doChat 方法
//            String chat = chatLanguageModel.chat(requestPrompt);// 正确的调用方式
            String chat = assistant.chat(content);
            parseAiAnswer(fileName, chat);
        } catch (Exception e) {
            extractPDFContent(filePath,fileName);
        }

    }

    public static void parseAiAnswer(String fileName, String json)  {

            // 按照指定分隔符分割字符串
            String[] split = json.split("```json");

            if (split.length > 1) {
                // 提取攻击技术部分
                Object o1 = JSON.parseObject(split[1].split("```")[0]);
                writeJsonToFile(fileName + "_tech.json", o1);
                Object o2 = JSON.parseObject(split[1].split("```")[0]);
                writeJsonToFile(fileName + "_plan.json", o2);
            }

    }
//
//    public static void parseAiAnswer(String fileName, String json) {
//
//        try{
//            // 按照指定分隔符分割字符串
//            String[] split = json.split("```json");
//
//            if (split.length > 1) {
//                // 提取攻击技术部分
//
//                Object o1 = JSON.parseObject(split[1].split("```")[0]);
//                writeJsonToFile(fileName + "_tech.json", o1);
//                Object o2 = JSON.parseObject(split[1].split("```")[0]);
//                writeJsonToFile(fileName + "_plan.json", o2);
//            }
//
//        }catch (Exception e){
//
//            HashMap<Object, Object> objectObjectHashMap = new HashMap<>();
//            objectObjectHashMap.put("json", json);
//            objectObjectHashMap.put("fileName", fileName);
//            objectObjectHashMap.put("error",e.getMessage());
//
//            writeJsonToFile(fileName + "_error.json",objectObjectHashMap);
//
//        }
//
//    }

    /**
     * 将 JSON 数据写入到指定目录的文件中
     *
     * @param fileName      文件名
     * @param content    要写入的 JSON 对象
     */
    public static void writeJsonToFile(String fileName, Object content) {
        String directoryPath = "output/json"; // 目标目录

        // 检查目录是否存在，如果不存在则创建
        File directory = new File(directoryPath);
        if (!directory.exists()) {
            directory.mkdirs(); // 创建目录
        }

        // 构建完整的文件路径
        String filePath = directoryPath + File.separator + fileName;

        // 将 JSON 对象转换为字符串

        // 写入文件
        try (BufferedWriter writer = new BufferedWriter(new FileWriter(filePath))) {
            writer.write(JSON.toJSONString(content));
            System.out.println("JSON数据已成功写入文件：" + filePath);
        } catch (IOException e) {
            System.err.println("写入文件时发生错误：" + e.getMessage());
            e.printStackTrace(); // 打印完整的堆栈信息
        }
    }


}
